import datetime
import os

import openpyxl
from dateutil.relativedelta import relativedelta
from langchain import PromptTemplate

from bot.openai_bot import OpenAIBot
from conf.config import BASE_DIR, logger
from prompts.prompt_templates import EXTRACT_CITY_NAME_PROMPT


def load_excel_sheet(excel_file_path, sheet_name):
    """
    加载excel某张表
    :param excel_file_path: excel文件的全路径
    :param sheet_name: 表名
    :return:
    """

    if not os.path.exists(excel_file_path):
        raise Exception(f'{excel_file_path}文件不存在，请根据data-template.xlsx模板创建{excel_file_path}文件')

    # 打开Excel文件（工作簿）
    book = openpyxl.load_workbook(excel_file_path)

    # 获取工作表
    sheet = book[sheet_name]

    # 将整张表转换为二维列表
    data = list(sheet.values)

    return data


def get_city_name(title: str, original_city_name: str = "") -> str:
    """
    从标题或直接从原本的城市名字段中获取城市名称
    :param title: 标题
    :param original_city_name: 原本的城市名
    :return: 城市名
    """
    if original_city_name:
        city_name = original_city_name
        logger.info(f'标题：{title} 原本的城市名：{city_name}')
        if city_name.endswith("市"):
            city_name = city_name[:-1]
    else:
        openai_bot = OpenAIBot(temperature=0)

        # 渲染模板
        prompt = PromptTemplate(
            input_variables=["title"],
            template=EXTRACT_CITY_NAME_PROMPT,
            template_format="jinja2"
        )
        prompt_formatted = prompt.format(title=title)

        city_name = openai_bot.reply(question=prompt_formatted)
        logger.info(f'标题：{title} 识别的城市名：{city_name}')

        if city_name in ("0", 0):
            return ""
        if city_name.endswith("。"):
            city_name = city_name[:-1]
        if city_name.endswith("市"):
            city_name = city_name[:-1]

    return city_name


def get_first_day_of_several_months_ago(months: int = 6):
    """
    获取几个月前的一号的日期
    :return:
    """

    today = datetime.date.today()
    several_months_ago = today - relativedelta(months=months)
    first_day_of_several_months_ago = datetime.date(several_months_ago.year, several_months_ago.month, 1).strftime("%Y-%m-%d")

    return first_day_of_several_months_ago


def load_title_list(excel_file_path: str = os.path.join(BASE_DIR, "data.xlsx")):
    """
    加载标题列表
    :param excel_file_path: 标题所在的excel文件路径
    :return:
    """
    sheet_name = "TitleList"
    data = load_excel_sheet(excel_file_path, sheet_name)

    title_list = list()
    for row in data[1:]:
        if not any(row):  # 空行
            continue
        title_list.append({"filename": row[0], "title": row[1], "city_name": row[3]})

    return title_list


def load_title_list_within_contexts(excel_file_path: str = os.path.join(BASE_DIR, "data.xlsx")):
    """
    加载标题列表含上下文
    :param excel_file_path: 标题所在的excel文件路径
    """
    sheet_name = "TitleList"
    data = load_excel_sheet(excel_file_path, sheet_name)

    title_list = list()
    for row in data[1:]:
        if not any(row):  # 空行
            continue
        title_list.append({"filename": row[0], "title": row[1], "contexts": row[2]})

    return title_list


def load_building_id_list(excel_file_path: str = os.path.join(BASE_DIR, "data.xlsx")):
    """
    加载楼盘id列表
    :param excel_file_path:  楼盘id所在的excel文件路径
    :return:
    """
    sheet_name = "BuildingList"
    data = load_excel_sheet(excel_file_path, sheet_name)

    building_id_list = list()
    for row in data[1:]:
        if not any(row):  # 空行
            continue
        building_id_list.append({"building_id": row[1], "building_name": row[2], "filename": f"{row[0]}"})

    return building_id_list


def main():
    ls = load_excel_sheet(os.path.join(BASE_DIR, "data.xlsx"), "Sheet1")

    # 按行打印
    for row in ls:
        if not any(row):  # 空行
            continue
        print(row)

    # # 按列打印
    # for row in ls:
    #     print(row[0])


if __name__ == '__main__':
    main()
